"""
multi_factor_macd.strategy

包含：指标计算、信号生成、仓位与成本模型辅助函数

假定输入 DataFrame 字段：['日期','股票代码','开盘','收盘','最高','最低','成交量','成交额','振幅','涨跌幅','涨跌额','换手率']
"""
from typing import Dict, Any, Tuple
import numpy as np
import pandas as pd


def ema(series: pd.Series, span: int) -> pd.Series:
    return series.ewm(span=span, adjust=False).mean()


def compute_indicators(df: pd.DataFrame, params: Dict[str, Any]) -> pd.DataFrame:
    df = df.copy()
    # ensure sorted by date
    df['日期'] = pd.to_datetime(df['日期'])
    df = df.sort_values('日期').reset_index(drop=True)

    close = df['收盘']
    high = df['最高']
    low = df['最低']
    vol = df['成交量']

    # MACD
    fast = params.get('macd_fast', 12)
    slow = params.get('macd_slow', 26)
    signal = params.get('macd_signal', 9)
    ema_fast = ema(close, fast)
    ema_slow = ema(close, slow)
    macd_line = ema_fast - ema_slow
    macd_signal = ema(macd_line, signal)
    macd_hist = macd_line - macd_signal

    df['macd_line'] = macd_line
    df['macd_signal'] = macd_signal
    df['macd_hist'] = macd_hist

    # RSI
    rsi_period = params.get('rsi_period', 14)
    delta = close.diff()
    up = delta.clip(lower=0)
    down = -delta.clip(upper=0)
    roll_up = up.ewm(alpha=1/rsi_period, adjust=False).mean()
    roll_down = down.ewm(alpha=1/rsi_period, adjust=False).mean()
    rs = roll_up / (roll_down.replace(0, np.nan))
    rsi = 100 - 100 / (1 + rs)
    df['rsi'] = rsi.fillna(50)

    # Moving averages
    ma_short = params.get('ma_short', 20)
    ma_long = params.get('ma_long', 50)
    df['ma_short'] = close.rolling(window=ma_short, min_periods=1).mean()
    df['ma_long'] = close.rolling(window=ma_long, min_periods=1).mean()

    # ATR
    atr_period = params.get('atr_period', 14)
    tr1 = high - low
    tr2 = (high - close.shift(1)).abs()
    tr3 = (low - close.shift(1)).abs()
    tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
    atr = tr.rolling(window=atr_period, min_periods=1).mean()
    df['atr'] = atr

    # volume normalized
    df['vol_z'] = (vol - vol.rolling(window=20, min_periods=1).mean()) / (vol.rolling(window=20, min_periods=1).std().replace(0, 1))

    return df


def generate_signals(df: pd.DataFrame, params: Dict[str, Any]) -> pd.DataFrame:
    """生成买卖信号，返回含 signals 列的 DataFrame

    规则（示例）:
    - MACD 主体为核心：hist 从负转正且大于阈值给+1分
    - RSI：低于 oversold 给多头加分，高于 overbought 给减分
    - 双均线：ma_short 上穿 ma_long 给加分，下穿给减分
    - 成交量过滤：成交量显著高于常态时放大信号

    将每项信号用权重合成得分，得分超过 entry_threshold 则买入；低于 exit_threshold 则卖出或平仓
    """
    df = df.copy()
    w_macd = params.get('w_macd', 0.4)
    w_rsi = params.get('w_rsi', 0.2)
    w_ma = params.get('w_ma', 0.2)
    w_vol = params.get('w_vol', 0.2)
    entry_th = params.get('entry_threshold', 0.5)
    exit_th = params.get('exit_threshold', 0.0)
    rsi_oversold = params.get('rsi_oversold', 30)
    rsi_overbought = params.get('rsi_overbought', 70)

    # signals normalized to [-1,1]
    df['sig_macd'] = np.tanh(df['macd_hist'] / (df['macd_line'].abs().rolling(20).mean().replace(0,1)))
    df['sig_rsi'] = 0.0
    df.loc[df['rsi'] < rsi_oversold, 'sig_rsi'] = 1.0
    df.loc[df['rsi'] > rsi_overbought, 'sig_rsi'] = -1.0
    df['sig_ma'] = 0.0
    df.loc[df['ma_short'] > df['ma_long'], 'sig_ma'] = 1.0
    df.loc[df['ma_short'] < df['ma_long'], 'sig_ma'] = -1.0
    df['sig_vol'] = np.tanh(df['vol_z'].fillna(0))

    df['composite_score'] = (w_macd * df['sig_macd'] + w_rsi * df['sig_rsi'] + w_ma * df['sig_ma'] + w_vol * df['sig_vol'])

    df['signal'] = 0
    # buy when composite_score > entry_th and macd histogram positive
    buy_mask = (df['composite_score'] > entry_th) & (df['macd_hist'] > 0)
    sell_mask = (df['composite_score'] < exit_th) & (df['macd_hist'] < 0)
    df.loc[buy_mask, 'signal'] = 1
    df.loc[sell_mask, 'signal'] = -1

    return df


def position_size(equity: float, price: float, atr: float, params: Dict[str, Any]) -> int:
    """基于ATR和风险限额计算交易手数（简化为股数）
    risk_per_trade: 单笔风险占比，如0.01
    stop_atr_mult: 止损设置为 n * ATR
    """
    risk_per_trade = params.get('risk_per_trade', 0.01)
    stop_atr_mult = params.get('stop_atr_mult', 2.0)
    stop_distance = stop_atr_mult * max(atr, 1e-6)
    # avoid zero division
    if stop_distance <= 0:
        return 0
    risk_amount = equity * risk_per_trade
    size = int(risk_amount / (stop_distance * price))
    return max(size, 0)


def apply_transaction_cost(price: float, size: int, side: str, volume: float, params: Dict[str, Any]) -> Tuple[float, float]:
    """返回 (effective_price, cost_amount)
    side: 'buy' or 'sell'
    滑点模型：base_slippage + coef / sqrt(volume)
    佣金与印花税按百分比
    """
    commission = params.get('commission', 0.0005)  # 0.05%
    stamp_tax = params.get('stamp_tax', 0.001)  # 0.1% (卖出)
    base_slip = params.get('base_slippage', 0.0002)
    slip_coef = params.get('slippage_coef', 1e5)

    # volume could be huge; ensure positive
    vol = max(volume, 1.0)
    slippage_pct = base_slip + slip_coef / (np.sqrt(vol) + 1e-9)
    if side == 'buy':
        eff_price = price * (1 + slippage_pct)
    else:
        eff_price = price * (1 - slippage_pct)

    trade_value = eff_price * size
    cost = trade_value * commission
    if side == 'sell':
        cost += trade_value * stamp_tax
    return eff_price, cost
